import os
import shutil
import pandas as pd
from collections import defaultdict


def fix_current_data_structure(data_path):
    """修复当前的数据结构问题"""

    print("=== 开始修复当前数据结构 ===")

    val_dir = os.path.join(data_path, 'val')

    # 1. 检查当前状态
    print("当前val目录结构:")
    for item in os.listdir(val_dir):
        item_path = os.path.join(val_dir, item)
        if os.path.isdir(item_path):
            subitems = os.listdir(item_path)
            print(f"  {item}/ -> {subitems[:3]}...")  # 显示前3个项目
        else:
            print(f"  {item}")

    # 2. 查找图片文件和标注文件
    all_images = []
    annotations_file = None

    # 遍历所有目录查找图片
    for root, dirs, files in os.walk(val_dir):
        for file in files:
            if file.lower().endswith(('.jpeg', '.jpg')):
                all_images.append(os.path.join(root, file))
            elif file == 'val_annotations.txt':
                annotations_file = os.path.join(root, file)

    print(f"找到图片文件: {len(all_images)} 个")
    print(f"找到标注文件: {annotations_file}")

    # 3. 读取标注文件
    if annotations_file and os.path.exists(annotations_file):
        df = pd.read_csv(annotations_file, sep='\t', header=None,
                         names=['filename', 'class', 'x', 'y', 'w', 'h'])
        print(f"标注文件记录数: {len(df)}")
    else:
        # 如果没有找到标注文件，尝试在val目录下找
        annotations_file = os.path.join(val_dir, 'val_annotations.txt')
        if os.path.exists(annotations_file):
            df = pd.read_csv(annotations_file, sep='\t', header=None,
                             names=['filename', 'class', 'x', 'y', 'w', 'h'])
            print(f"标注文件记录数: {len(df)}")
        else:
            print("❌ 找不到标注文件，无法继续修复")
            return False

    # 4. 按类别重新组织数据
    class_groups = defaultdict(list)
    for _, row in df.iterrows():
        class_groups[row['class']].append(row)

    print(f"需要创建的类别数: {len(class_groups)}")

    # 5. 创建正确的目录结构
    for class_name, rows in class_groups.items():
        class_dir = os.path.join(val_dir, class_name)
        class_images_dir = os.path.join(class_dir, 'images')

        # 创建目录
        os.makedirs(class_images_dir, exist_ok=True)

        # 移动图片文件
        moved_count = 0
        for row in rows:
            # 查找图片文件的实际位置
            img_found = False
            for img_path in all_images:
                if os.path.basename(img_path) == row['filename']:
                    # 移动文件到正确位置
                    dst_path = os.path.join(class_images_dir, row['filename'])
                    if not os.path.exists(dst_path):  # 避免重复移动
                        shutil.move(img_path, dst_path)
                    moved_count += 1
                    img_found = True
                    break

            if not img_found:
                print(f"⚠️ 未找到图片: {row['filename']}")

        # 创建该类别的标注文件
        class_annotations = os.path.join(class_dir, 'val_annotations.txt')
        class_df = pd.DataFrame(rows)
        class_df.to_csv(class_annotations, sep='\t', index=False, header=False)

        print(f"✅ 类别 {class_name}: 移动了 {moved_count}/{len(rows)} 张图片")

    # 6. 清理空目录
    print("\n=== 清理空目录 ===")
    for item in os.listdir(val_dir):
        item_path = os.path.join(val_dir, item)
        if os.path.isdir(item_path) and item not in class_groups.keys():
            # 检查是否是空目录
            try:
                if len(os.listdir(item_path)) == 0:
                    shutil.rmtree(item_path)
                    print(f"删除空目录: {item}")
                else:
                    print(f"保留非空目录: {item}")
            except:
                pass

    # 7. 验证修复结果
    print("\n=== 修复结果验证 ===")
    total_images = 0
    for class_name in list(class_groups.keys())[:5]:  # 检查前5个类别
        class_images_dir = os.path.join(val_dir, class_name, 'images')
        if os.path.exists(class_images_dir):
            images = [f for f in os.listdir(class_images_dir)
                      if f.lower().endswith(('.jpeg', '.jpg'))]
            total_images += len(images)
            print(f"类别 {class_name}: {len(images)} 张图片")

    print(f"验证集总图片数: {total_images}")
    print("✅ 数据结构修复完成！")

    return True


# 运行修复
data_path = "tiny-imagenet-200"
success = fix_current_data_structure(data_path)